package ws

import (
	"net"
	"sync"
	"errors"
	"encoding/base64"

	"golang.org/x/net/websocket"
)

type webSocketConn struct {
	sync.Mutex
	conn 		*websocket.Conn
	isClosed	bool
	readBuf		[]byte
}

func NewWebSocketConn(conn *websocket.Conn) *webSocketConn {
	c := &webSocketConn{}
	c.conn = conn
	c.isClosed = false
	c.readBuf = make([]byte, 1024)
	return c
}

func (conn *webSocketConn) WriteMsg(args ...[]byte)  (int, error) {
	conn.Lock()
	defer conn.Unlock()
	if !conn.isClosed {
		msgLen := 0
		for i := 0; i < len(args); i++ {
			msgLen += len(args[i])
		}
		if msgLen > 0 {
			msg := make([]byte, msgLen)
			lenght := 0
			for i := 0; i < len(args); i++ {
				copy(msg[lenght:], args[i])
				lenght += len(args[i])
			}

			buf := make([]byte, base64.StdEncoding.EncodedLen(len(msg)))
			base64.StdEncoding.Encode(buf, msg)
			return conn.conn.Write(buf)
		}
	}
	return 0, errors.New("conn is closed")
}
	
func (conn *webSocketConn) WriteWithBase64(buf []byte, buf64 []byte)  (int, error) {
	conn.Lock()
	defer conn.Unlock()
	if !conn.isClosed {
		return conn.conn.Write(buf64)
	}
	return 0, errors.New("conn is closed")
}

func (conn *webSocketConn) ReadMsg()([]byte, error) {
	n, err := conn.conn.Read(conn.readBuf)
	if nil == err && n > 0 {
		rb, err := base64.StdEncoding.DecodeString(string(conn.readBuf[:n]))
		return rb, err
	}
	return nil, err
}

func (conn *webSocketConn) LocalAddr() net.Addr {
	return conn.conn.LocalAddr()
}

func (conn *webSocketConn) RemoteAddr() net.Addr {
	return conn.conn.RemoteAddr()
}

func (conn *webSocketConn) Close() {
	conn.Lock()
	defer conn.Unlock()
	if !conn.isClosed {
		// conn.conn.SetReadDeadline(time.Now().Add(10 * time.Second))
		conn.conn.Close()
		conn.isClosed = true
	}
}
